import datetime
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pylab
import talib as ta
import talib as tl
from hs_udata import set_token
from matplotlib import dates as mdates
from matplotlib import ticker as mticker
from mplfinance.original_flavor import candlestick_ohlc

warnings.filterwarnings("ignore")


# 上穿信号
# values上穿avg，返回上穿位置索引
def upCross(values, window):
    avg = tl.MA(values, window)
    assert len(values) == len(avg), '上穿信号输入维度不相等'
    assert len(values) > 1, '上穿信号长度至少为2'
    res = []
    for i in range(window, len(values)):
        if values[i - 1] < avg[i - 1] and values[i] > avg[i]:
            val = [values[i], i]
            res.append(val)

    return res


# 下穿信号
# values下穿avg，返回下穿位置索引
def downCross(values, window):
    avg = tl.MA(values, window)
    assert len(values) == len(avg), '上穿信号输入维度不相等'
    assert len(values) > 1, '上穿信号长度至少为2'
    res = []
    for i in range(window, len(values)):
        if values[i - 1] > avg[i - 1] and values[i] < avg[i]:
            val = [values[i], i]
            res.append(val)

    return res


def movingaverage(values, window):
    weigths = np.repeat(1.0, window) / window
    smas = np.convolve(values, weigths, 'valid')
    return smas  # as a numpy array


def RSI(array_list, periods=14):
    length = len(array_list)
    rsies = [np.nan] * length

    if length < periods or length == periods:  # 判断传入的数据 条(天) 是否小于等于 14条(天) 直接返回rsies
        return rsies
    up_avg = 0
    down_avg = 0

    first_t = array_list[:periods + 1]  # 从第0条开始取15条
    for i in range(1, len(first_t)):  # 遍历 first_t
        if first_t[i] > first_t[i - 1] or first_t[i] == first_t[i - 1]:  # 数据 后者>=前者
            up_avg += first_t[i] - first_t[i - 1]  # 数据  计算前后差值 放入up_avg 上平均
        else:
            down_avg += first_t[i - 1] - first_t[i]  # 数据  计算前后差值 放入down_avg 下平均
    up_avg = up_avg / periods  # 重新计算平均值
    down_avg = down_avg / periods  # 重新计算平均值
    rs = up_avg / down_avg  # 计算 上平均 / 下平均
    rsies[periods] = 100 - 100 / (1 + rs)  # 不知道啥啥公式

    for j in range(periods + 1, length):  # 遍历传入数据 前15个 到 之后的值
        up = 0
        down = 0
        if array_list[j] > array_list[j - 1] or array_list[j] == array_list[j - 1]:  # 数据 后者>=前者
            up = array_list[j] - array_list[j - 1]  # 前后者差值
            down = 0
        else:
            up = 0
            down = array_list[j - 1] - array_list[j]  # 前后者差值
        up_avg = (up_avg * (periods - 1) + up) / periods  # 上平均 = (上平均/13 + 上差值)/ 14
        down_avg = (down_avg * (periods - 1) + down) / periods  # 下平均 = (下平均/13 + 下差值)/ 14
        rs = up_avg / down_avg  # 计算 上平均 / 下平均
        rsies[j] = 100 - 100 / (1 + rs)  # 不知道啥啥公式
    return rsies


def loadImg(data, MA1, MA2, num, cross):
    data = data.loc[data.turnover_status == '交易']  # 剔除非交易日

    dataReshape = data[['trading_date', 'open_price', 'high_price', 'low_price', 'close_price', 'business_amount']]
    dataReshape[['open_price', 'high_price', 'low_price', 'close_price', 'business_amount']] = dataReshape[
        ['open_price', 'high_price', 'low_price', 'close_price', 'business_amount']].astype(float)  # 将价格数据类型转为浮点数

    dateCopy = dataReshape['trading_date']  # 备份时间刻度
    dataReshape['trading_date'] = dataReshape.index  # 将时间刻度替换成索引 以免k线计算 留空白
    X_axis = dataReshape.index[-num:]  # 统一使用索引作为刻度

    # 画K线图------------------------------------------------------

    ax1 = plt.subplot2grid((6, 4), (1, 0), rowspan=4, colspan=4, facecolor='#07000d')

    resup = upCross(dataReshape.close_price.values[-num:], cross)
    for r in resup:
        ax1.plot(X_axis[r[1]], r[0], 'o', color='lightcoral', alpha=.8)

    resdown = downCross(dataReshape.close_price.values[-num:], cross)
    for r in resdown:
        ax1.plot(X_axis[r[1]], r[0], 'o', color='limegreen', alpha=.8)

    candlestick_ohlc(
        ax1, dataReshape.values[-num:], width=.5, colorup='#ff1717', colordown='#53c156')

    Av1 = movingaverage(dataReshape.close_price.values, MA1)
    Av2 = movingaverage(dataReshape.close_price.values, MA2)
    Label1 = str(MA1) + ' SMA'
    Label2 = str(MA2) + ' SMA'

    ax1.plot(X_axis, Av1[-num:],
             '#e1edf9', label=Label1, linewidth=1.5)
    ax1.plot(X_axis, Av2[-num:],
             '#4ee6fd', label=Label2, linewidth=1.5)

    ax1.grid(True, color='w')  # 启用网格

    ax1.yaxis.label.set_color("w")
    ax1.xaxis.set_major_locator(mticker.MaxNLocator(10))  # 调整纵向格数
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))  # 将日期可视化

    ax1.spines['bottom'].set_color("#5998ff")
    ax1.spines['top'].set_color("#5998ff")
    ax1.spines['left'].set_color("#5998ff")
    ax1.spines['right'].set_color("#5998ff")
    ax1.tick_params(axis='y', colors='w')  # 改变刻度、刻度标签、网格线的外观
    plt.gca().yaxis.set_major_locator(mticker.MaxNLocator(prune='upper'))  # 调整横向格数
    ax1.tick_params(axis='x', colors='w')
    # ---------------------------------------------------------------

    # 成交量图----------------------------------------------------------
    volumeMin = 0
    ax1v = ax1.twinx()  # 获取共享X轴的第二个X轴
    ax1v.fill_between(X_axis, volumeMin, dataReshape.business_amount.values[-num:],
                      facecolor='#00ffe8', alpha=.4)  # alpha：覆盖区域的透明度[0,1],其值越大，表示越不透明
    # ax1v.axes.yaxis.set_ticklabels([]) #隐藏刻度值
    ax1v.grid(False)  # 关闭网格

    # 将它编辑为3，这样它会大一些
    ax1v.set_ylim(0, 3 * dataReshape.business_amount.values.max())

    ax1v.spines['bottom'].set_color("#5998ff")
    ax1v.spines['top'].set_color("#5998ff")
    ax1v.spines['left'].set_color("#5998ff")
    ax1v.spines['right'].set_color("#5998ff")
    ax1v.tick_params(axis='x', colors='w')
    ax1v.tick_params(axis='y', colors='w')
    plt.ylabel('Stock price and business_amount')
    # -------------------------------------------------------------------

    # 画RSI-------------------------------------------------------------
    # maLeg = plt.legend(loc=9, ncol=2, prop={'size': 7}, fancybox=False, borderaxespad=0)
    # maLeg.get_frame().set_alpha(0.4)
    # textEd = pylab.gca().get_legend().get_texts()
    # pylab.setp(textEd[0:5], color=
    # 'w')

    ax0 = plt.subplot2grid((6, 4), (0, 0), sharex=ax1, rowspan=1, colspan=4, facecolor='#07000d')
    rsi = np.array(RSI(dataReshape.close_price.values))

    rsiCol = '#c1f9f7'
    posCol = '#386d13'
    negCol = '#8f2020'

    ax0.plot(X_axis, rsi[-num:], rsiCol, linewidth=1.5)
    ax0.axhline(70, color=negCol)  # 绘制参考系
    ax0.axhline(30, color=posCol)  # 绘制参考系

    ax0.fill_between(X_axis, rsi[-num:], 70, where=(  # 填色
            rsi[-num:] >= 70), facecolor=negCol, edgecolor=negCol, alpha=0.5)

    ax0.fill_between(X_axis, rsi[-num:], 30, where=(  # 填色
            rsi[-num:] <= 30), facecolor=posCol, edgecolor=posCol, alpha=0.5)

    ax0.set_yticks([30, 70])  # 设置带有刻度列表的y刻度
    ax0.yaxis.label.set_color("w")

    ax0.spines['bottom'].set_color("#5998ff")  # 边框
    ax0.spines['top'].set_color("#5998ff")  # 边框
    ax0.spines['left'].set_color("#5998ff")  # 边框
    ax0.spines['right'].set_color("#5998ff")  # 边框
    ax0.tick_params(axis='y', colors='w')
    ax0.tick_params(axis='x', colors='w')
    plt.ylabel('RSI')
    # -------------------------------------------------------------------

    # 画macd---------------------------------------------------------------
    ax2 = plt.subplot2grid((6, 4), (5, 0), sharex=ax1,
                           rowspan=1, colspan=4, facecolor='#07000d')
    fillcolor = '#00ffe8'
    nema = 9
    emaslow, emafast, macd = ta.MACD(dataReshape.close_price.values)

    ema9 = ta.EMA(macd, nema)
    ax2.plot(X_axis, macd[-num:], color='#4ee6fd', lw=2)
    ax2.plot(X_axis, ema9[-num:], color='#e1edf9', lw=1)
    ax2.fill_between(X_axis, macd[-num:] - ema9[-num:],  # 填色
                     0, alpha=0.5, facecolor=fillcolor, edgecolor=fillcolor)
    plt.gca().yaxis.set_major_locator(mticker.MaxNLocator(prune='upper'))  # 调整横向格数

    ax2.spines['bottom'].set_color("#5998ff")
    ax2.spines['top'].set_color("#5998ff")
    ax2.spines['left'].set_color("#5998ff")
    ax2.spines['right'].set_color("#5998ff")
    ax2.tick_params(axis='x', colors='w')
    ax2.tick_params(axis='y', colors='w')
    plt.ylabel('MACD', color='w')
    ax2.yaxis.set_major_locator(mticker.MaxNLocator(nbins=5, prune='upper'))  # 调整横向格数 y刻度最大数

    for label in ax2.xaxis.get_ticklabels():  # 获取日期刻度标签
        label.set_rotation(0)  # 旋转

    plt.suptitle(data.prod_code.values[-1], color='w')
    plt.setp(ax0.get_xticklabels(), visible=False)
    plt.setp(ax1.get_xticklabels(), visible=False)

    plt.xticks(X_axis, dateCopy[-num:], rotation=0)
    plt.gca().xaxis.set_major_locator(mticker.MaxNLocator(10))


data = pd.read_csv("test.csv", parse_dates=True, encoding='UTF-8')
undate = data.trading_date.values[-1]
undate = datetime.datetime.strptime(undate, '%Y-%m-%d')
set_token(token='ruXPpxAnX4-5DOF7SXgKOuI4y7WOtBPVXqE79Nj2q30dHluNv6Hd-ylE-JwynvBZ')  # 注册后，获取并替换token

# plt.ion()
fig = plt.figure(facecolor='#07000d', figsize=(15, 10))

loadImg(data, 5, 30, 50, 30)

# fig.canvas.draw()
# fig.canvas.flush_events()

plt.show()

# loadImg(数据,MA1,MA2,展示数目,plt.figure())


# class SituationMap:
#     def __init__(this):

#     def LoadSituationMap(data,MK1,MK2,num,cross):
#         loadImg(data,MK1,MK2,num,cross)
